#include "DFT.h"
using namespace ypl;

/************************
 * 采用复数数组的接口
 ************************/

/*!
 * @brief 一维（逆）离散傅里叶变换（复数输入）
 */
void DFT::operator ()(const Array1D<Complex> &aIn,
                      Array1D<Complex> &aOut,
                      Direct direct)
{
    Array1D<SComplex> asIn, asOut;
    complexToSComplex(aIn, asIn);
    complexToSComplex(aOut, asOut);
    this->operator()(asIn, asOut, direct); ///< 调用简单复数接口
    scomplexToComplex(asOut, aOut);
}

/*!
 * @brief 一维离散傅里叶变换（实数输入）
 */
void DFT::operator ()(const Array1D<Real> &aIn,
                      Array1D<Complex> &aOut,
                      Direct direct)
{
    Array1D<SComplex> asOut;
    complexToSComplex(aOut, asOut);
    this->operator()(aIn, asOut, direct); ///< 调用简单复数接口
    scomplexToComplex(asOut, aOut);
}

/*!
 * @brief 二维离散傅里叶变换（复数输入）
 */
void DFT::operator ()(Array2D<Complex> &aIn,
                      Array2D<Complex> &aOut,
                      Direct direct)
{
    Array2D<SComplex> asIn, asOut;
    complexToSComplex(aIn, asIn);
    complexToSComplex(aOut, asOut);
    this->operator()(asIn, asOut, direct); ///< 调用简单复数接口
    scomplexToComplex(asOut, aOut);
}

/*!
 * @brief 二维离散傅里叶变换（实数输入）
 */
void DFT::operator ()(const Array2D<Real> &aIn,
                      Array2D<Complex> &aOut,
                      Direct direct)
{
    Array2D<SComplex> asOut;
    complexToSComplex(aOut, asOut);
    this->operator()(aIn, asOut, direct); ///< 调用简单复数接口
    scomplexToComplex(asOut, aOut);
}


/***************************
 * 采用简单复数数组的接口
 ***************************/


/*!
 * @brief 一维（逆）离散傅里叶变换（复数输入）
 */
void DFT::operator ()(const Array1D<SComplex> &aIn,
                      Array1D<SComplex> &aOut,
                      Direct direct)
{
    /*! 转换参数以适应底层计算接口 */
    Integral n = aIn.dim();                ///< 待变换的复数个数
    if (n == 0)
    {
        aOut = Array1D<SComplex>(0);
        return;
    }
    if (aOut.dim() != n)
    {
        aOut = Array1D<SComplex>(n);
    }
    const SComplex *pIn = aIn;
    SComplex *pOut = aOut;

    /*! 调用底层算法计算（逆）离散傅里叶变换 */

    /*! 点数为 2 的幂形式，则用快速傅里叶变换，否则直接计算 */
    if (isPower2(n))
    {
        fft(pIn, n, pOut, direct);
    }
    else
    {
        dft(pIn, n, pOut, direct);
    }

}

/*!
 * @brief 一维离散傅里叶变换（实数输入）
 */
void DFT::operator ()(const Array1D<Real> &aIn,
                      Array1D<SComplex> &aOut,
                      Direct direct)
{
    /*! 实数数组转化为简单复数数组 */
    Integral dim = aIn.dim();
    Array1D<SComplex> asIn(dim);
    for (Integral i = 0; i < dim; ++i)
    {
        asIn[i].real = aIn[i];
        asIn[i].imag = 0.0;
    }
    this->operator()(asIn, aOut, direct); ///< 调用简单复数数组输入形式
}

/*!
 * @brief 二维离散傅里叶变换（简单复数输入）
 */
void DFT::operator ()(Array2D<SComplex> &aIn,
                      Array2D<SComplex> &aOut,
                      Direct direct)
{
    /*! 转换参数以适应底层计算接口 */
    Integral dim1 = aIn.dim1();
    Integral dim2 = aIn.dim2();
    if (aOut.dim1() != dim1 || aOut.dim2() != dim2)
    {
        aOut = Array2D<SComplex>(dim1, dim2);
    }
    SComplex **pIn = aIn;
    SComplex **pOut = aOut;

    /*! 调用底层算法计算（逆）离散傅里叶变换 */

    /*! 检查点数是否为 2 的幂形式 */
    bool isPower2 = false;
    if ( (dim1 & (dim1-1)) == 0 && (dim2 & (dim2-1)) == 0 )
    {
        isPower2 = true;
    }
    /*! 点数为 2 的幂形式，则用快速傅里叶变换，否则直接计算 */
    if (isPower2)
    {
        fft(pIn, dim1, dim2, pOut, direct);
    }
    else
    {
        dft(pIn, dim1, dim2, pOut, direct);
    }

}

/*!
 * @brief 二维离散傅里叶变换（实数输入）
 */
void DFT::operator ()(const Array2D<Real> &aIn,
                      Array2D<SComplex> &aOut,
                      Direct direct)
{
    /*! 实数数组转化为简单复数数组 */
    Integral dim1 = aIn.dim1();
    Integral dim2 = aIn.dim2();
    Array2D<SComplex> asIn(dim1, dim2);
    for (Integral i = 0; i < dim1; ++i)
    {
        for (Integral j = 0; j < dim2; ++j)
        {
            asIn[i][j].real = aIn[i][j];
            asIn[i][j].imag = 0.0;
        }

    }
    this->operator()(asIn, aOut, direct); ///< 调用简单复数数组输入形式
}


/*!
 * @brief 复数序列的一维离散傅里叶变换，直接计算
 * @param pIn       -in- 复数序列
 * @param n         -in- 数据点数（任意正整数）
 * @param pOut      -out- 离散傅里叶变换结果序列
 * @param direct    -in- 傅里叶变换的方向，
 *                       可选: Forward（正变换，默认）, Backward（逆变换）,
 *
 */
void DFT::dft(const SComplex *pIn,
              Integral n,
              SComplex *pOut,
              Direct direct)
{
    Array1D<SComplex> aW;
    calcW(aW, n, n, direct);
    SComplex *w = aW;
    for (Integral k = 0; k < n; ++k)
    {
        pOut[k].real = 0.0;
        pOut[k].imag = 0.0;
        for (Integral j = 0; j < n; ++j)
        {
            Integral kj = k*j;
            if (kj >= n) kj %= n;
            pOut[k].real += pIn[j].real * w[kj].real
                            - pIn[j].imag * w[kj].imag;
            pOut[k].imag += pIn[j].real * w[kj].imag
                            + pIn[j].imag * w[kj].real;
        }
    }
    if (direct == Backward)
    {
        divide(pOut, n, n);
    }
}



/*!
 * @brief 复数序列的一维快速傅里叶变换
 * @param aIn -in- 复数序列，复数个数必须为2的幂,不小于2
 * @param n   -in- 数据点数（任意正整数）
 * @param aOut -out- 离散傅里叶变换结果序列
 * @param direct -in- 傅里叶变换的方向，
 *        可选为 Forward（正变换）, Backward（逆变换）,
 *        默认为 Forward
 */
void DFT::fft(const SComplex *pIn,
              Integral n,
              SComplex *pOut,
              Direct direct)
{
    /*! 按位反转复制 */
    Array1D<Integral> newOrder;
    bitReverse(n, newOrder);            ///< 按位反转
    orderCopy(pIn, n, pOut, newOrder);  ///< 复制到简单复数数组

    /*! 计算需要用到的单位复根 w[0..n/2-1] */
    Integral n2 = n >> 1;           ///< n2 = n/2
    Array1D<SComplex> aW(n2);       ///< 此处用 w[] 会出错, 因为原生数组在栈上分配存储，容量小
    calcW(aW, n, n2, direct);
    SComplex *w = aW;

    calcKernel(pOut, n, w);

    /*! 若为逆变换，则除以点数 */
    if (direct == Backward)
    {
        divide(pOut, n, n);
    }
}


/*!
 * @brief 复数序列的二维离散傅里叶变换，直接计算
 * @param pIn       -in- 复数序列
 * @param dim1      -in- 数据第一维点数（任意正整数）
 * @param dim2      -in- 数据第二维点数（任意正整数）
 * @param pOut      -out- 离散傅里叶变换结果序列
 * @param direct    -in- 傅里叶变换的方向，
 *                       可选: Forward（正变换，默认）, Backward（逆变换）,
 *
 */
void DFT::dft(SComplex **pIn,
              Integral dim1,
              Integral dim2,
              SComplex **pOut,
              Direct direct)
{

}



/*!
 * @brief 复数序列的二维快速傅里叶变换
 * @param aIn       -in- 复数序列，复数行列数必须均为 2 的幂
 * @param dim1      -in- 数据第一维点数（必须为 2 的幂）
 * @param dim2      -in- 数据第二维点数（必须为 2 的幂）
 * @param aOut      -out- 离散傅里叶变换结果序列
 * @param direct    -in- 傅里叶变换的方向，
 *        可选为 Forward（正变换）, Backward（逆变换）,
 *        默认为 Forward
 */
void DFT::fft(SComplex **pIn,
              Integral dim1,
              Integral dim2,
              SComplex **pOut,
              Direct direct)
{

    /* 对各行进行 bit 反转复制 */
    Array1D<Integral> newOrder;
    bitReverse(dim2, newOrder);
    for (Integral iRow = 0; iRow < dim1; ++iRow)
    {
        orderCopy(pIn[iRow], dim2, pOut[iRow], newOrder);
    }


    /*!
     * 计算行 dft 需要用到的单位复根
     *    wn[0..m-1],
     * n = dim2, m = dim2 / 2
     */
    Integral m = dim2 >> 1;
    Array1D<SComplex> aW(m);
    calcW(aW, dim1, m, direct);
    SComplex *pW = aW;

    /*! 计算各行的 dft */
    calcKernelDim1(pOut, dim1, dim2, pW);

    /*! 对各列进行 bit 反转复制 */
    if( dim1 != dim2)
    {
        bitReverse(dim1, newOrder);
        m = dim1 >> 1; ///< m = dim1 / 2;

        /*! 计算各列 dft 需要用到的单位复根 wn[0..m-1], n=dim1 */
        calcW(aW, dim1, m, direct);
    }
    Integral *pNewOrder = newOrder;

    for (Integral i = 0; i < dim1; ++i)
    {
        Integral iNew = pNewOrder[i];
        for (Integral iColumn = 0; iColumn < dim2; ++iColumn)
        {
            if (i < iNew)
            {
                SComplex ctmp = pOut[i][iColumn];
                pOut[i][iColumn] = pOut[iNew][iColumn];
                pOut[iNew][iColumn] = ctmp;
            }

        }
    }

    pW = aW;

    /*! 计算各列的 dft */
    calcKernelDim2(pOut, dim1, dim2, pW);

    /*! 若为逆变换，则除以点数 */
    if (direct == Backward)
    {
        divide(pOut, dim1, dim2, dim1*dim2);
    }
}


/*!
 * @brief 按位反转
 * @param n     -in-    指定对 0~n-1 进行反转，n 必须为 2的幂（2^k）
 * @param a     -out-   保存 0~n-1 反转后的顺序
 */
void DFT::bitReverse(Integral n, Array1D<Integral> &a)
{
    if (a.dim() != n)
    {
        a = Array1D<Integral>(n);
    }
    Integral *pa = a; ///< 用指针访问，提高效率
    pa[0] = 0;
    Integral n2 = n >> 1;    ///< n2 = n / 2
    for (Integral i = 1; i < n; ++i)
    {
        Integral m = n2;
        pa[i] = pa[i-1];
        while ( pa[i] >= m)
        {
            pa[i] -= m;
            m >>= 1;
        }
        pa[i] += m;
    }
}

/*!
 * @brief
   计算单位复根 w[0..m-1]，w[k] = w[k-1] * wn
       wn = exp(-2*pi*i/n) （正变换）
       wn = exp(2*pi*i/n)  （逆变换）
   迭代计算，采用
       w[k] = w[k-1] + w[k-1] * ( wn -1 )
   以减小误差积累 ( n^(1/2) )
       wn1 = wn - 1 = cos(theta)-1  + i*sin(theta)
*/
void DFT::calcW(Array1D<SComplex> &aW, Integral n, Integral m, Direct direct)
{
    Real pi_2 = 2.0 * Constants::pi;
    if (direct == Forward)      // 正变换添加负号
    {
        pi_2 = -pi_2;
    }
    Real theta = pi_2 / n;
    Real tmp = sin(0.5 * theta);
    SComplex wn1;                  // 逆 wn1 = wn - 1
    wn1.real = -2 * tmp * tmp;    // real = cos(theta) - 1 = -2*sin(theta/2)*sin(theta/2)
    wn1.imag = sin(theta);

    if (m == 0 || m > n)
    {
        m = n;
    }
    if (aW.dim() != m)
    {
        aW = Array1D<SComplex>(m);
    }
    SComplex *w = aW;
    w[0].real = 1.0;
    w[0].imag = 0.0;
    /* w[j] = w[j-1] * wn = w[j-1] + w[j-1] * wn1 */
    for (Integral j = 1; j < m; ++j)
    {
        SComplex &wj1 = w[j-1];
        w[j].real = wj1.real + wj1.real * wn1.real
                    - wj1.imag * wn1.imag;
        w[j].imag = wj1.imag + wj1.real * wn1.imag
                    + wj1.imag * wn1.real;
    }
}


/*!
 * @brief fft 计算核心
 *        计算给定一维复数序列的按位反序序列的离散傅里叶变换（逆变换不除 n）
 * @param pData     -in-    复数序列
 * @param n         -in-    复数序列长度
 * @param pW        -in-    单位根序列 pW[k]=wn^k, k 至少取 0..n/2-1
 */
void DFT::calcKernel(SComplex *pData, Integral n, const SComplex *pW)
{
    /*! 计算幂次数 s: n = 2^s, s = log(n) */
    Integral s = log2(n);

    for (Integral i = 1; i <= s; ++i)
    {
        Integral m = 1 << i;           ///< m = 2^i, m=2,4,..., dim;
        Integral m2 = m >> 1;          ///< m2 = m/2;
        Integral si = s - i;
        /*!
         *  循环次数多的放在内层效果更好
         *  虽然调换两for循环顺序，可以减少 jj 的计算
         */
        for (Integral k = 0; k < n; k += m)
        {
            for (Integral j = 0; j < m2; ++j)
            {
                Integral kj = k + j;
                Integral jj = j << si;  // jj = j * 2^(s-i), 满足 wm(j) = wn(jj)
                Integral kjm2 = kj + m2;
                SComplex &cLow = pData[kj];
                SComplex &cHigh = pData[kjm2];
                const SComplex &w = pW[jj];
                SComplex t;         // t = w * cHigh
                t.real = w.real * cHigh.real
                         - w.imag * cHigh.imag;
                t.imag = w.real * cHigh.imag
                         + w.imag * cHigh.real;
                cHigh.real = cLow.real - t.real;
                cHigh.imag = cLow.imag - t.imag;
                cLow.real += t.real;
                cLow.imag += t.imag;
            }
        }
    }
}


/*!
 * @brief fft 计算核心
 *        计算给定二维复数序列【各行】
 *        的按位反序序列的离散傅里叶变换（逆变换不正则化））
 * @param pData     -in-    复数序列
 * @param n         -in-    复数序列长度
 * @param pW        -in-    单位根序列 pW[k]=wn^k, k 至少取 0..n/2-1
 */
void DFT::calcKernelDim1(SComplex **pData,
                         Integral dim1,
                         Integral dim2,
                         const SComplex *pW)
{
    /*! 幂次数 s: n = 2^s, s = log(n) */
    Integral s = log2(dim2);

    for (Integral i = 1; i <= s; ++i)
    {
        Integral m = 1 << i;           ///< m = 2^i, m=2,4,..., dim;
        Integral m2 = m >> 1;          ///< m2 = m/2;
        Integral si = s - i;
        for (Integral k = 0; k < dim2; k += m)
        {
            for (Integral j = 0; j < m2; ++j)
            {
                Integral kj = k + j;
                Integral jj = j << si;  ///< jj = j * 2^(s-i), 满足 wm(j) = wn(jj)
                Integral kjm2 = kj + m2;
                for (Integral iRow = 0; iRow < dim1; ++iRow)
                {
                    SComplex &cLow = pData[iRow][kj];
                    SComplex &cHigh = pData[iRow][kjm2];
                    const SComplex &w = pW[jj];
                    SComplex t;         ///< t = w * cHigh
                    t.real = w.real * cHigh.real
                             - w.imag * cHigh.imag;
                    t.imag = w.real * cHigh.imag
                             + w.imag * cHigh.real;
                    cHigh.real = cLow.real - t.real;
                    cHigh.imag = cLow.imag - t.imag;
                    cLow.real += t.real;
                    cLow.imag += t.imag;
                }
            }
        }
    }
}

/*!
 * @brief fft 计算核心
 *        计算给定二维复数序列【各列】
 *        的按位反序序列的离散傅里叶变换（逆变换不正则化）
 * @param pData     -in-    复数数据
 * @param dim1      -in-    复数数据行数
 * @param dim2      -in-    复数数据列数
 * @param pW        -in-    单位根序列 pW[k]=wn^k, k 至少取 0..n/2-1
 */
void DFT::calcKernelDim2(SComplex **pData,
                         Integral dim1,
                         Integral dim2,
                         const SComplex *pW)
{
    /*! 幂次数 s: n = 2^s, s = log(n) */
    Integral s = log2(dim1);

    for (Integral i = 1; i <= s; ++i)
    {
        Integral m = 1 << i;           ///< m = 2^i, m=2,4,..., dim;
        Integral m2 = m >> 1;          ///< m2 = m/2;
        Integral si = s - i;
        for (Integral k = 0; k < dim1; k += m)
        {
            for (Integral j = 0; j < m2; ++j)
            {
                Integral kj = k + j;
                Integral jj = j << si;  ///< jj = j * 2^(s-i), 满足 wm(j) = wn(jj)
                Integral kjm2 = kj + m2;
                for (Integral iColumn = 0; iColumn < dim2; ++iColumn)
                {
                    SComplex &cHigh = pData[kjm2][iColumn];
                    SComplex &cLow = pData[kj][iColumn];
                    const SComplex &w = pW[jj];
                    SComplex t;         ///< t = w * cHigh
                    t.real = w.real * cHigh.real
                             - w.imag * cHigh.imag;
                    t.imag = w.real * cHigh.imag
                             + w.imag * cHigh.real;
                    cHigh.real = cLow.real - t.real;
                    cHigh.imag = cLow.imag - t.imag;
                    cLow.real += t.real;
                    cLow.imag += t.imag;
                }
            }
        }
    }
}

void DFT::divide(SComplex *pData, Integral n, Integral divisor)
{
    if (divisor == 0)
    {
        divisor = n;
    }
    Real multiplier = 1.0 / divisor;
    for (Integral i = 0; i < n; ++i)
    {
        pData[i].real *= multiplier;
        pData[i].imag *= multiplier;
    }
}

void DFT::divide(SComplex **pData, Integral dim1, Integral dim2, Integral divisor)
{
    if (divisor == 0)
    {
        divisor = dim1*dim2;
    }
    Real multiplier = 1.0 / divisor;
    for (Integral i = 0; i < dim1; ++i)
    {
        for (Integral j = 0; j < dim2; ++j)
        {
            pData[i][j].real *= multiplier;
            pData[i][j].imag *= multiplier;
        }
    }
}

/*! @brief 计算 k = [log2(n)]，
 *         log2(n) 为满足 2^k <= n 的最大整数 k
 *         若 n=2^k，则 log2(n) = k
 */
Integral DFT::log2(Integral n)
{
    /* 计算幂次数 s: n = 2^s, s = log(n) */
    Integral s = 0;
    Integral ntmp = n >> 1;
    while (ntmp > 0)
    {
        ++s;
        ntmp >>= 1; // ntmp /= 2
    }
    return s;
}

/*! @brief 检查 n 是否为 2^k 形式 */
bool DFT::isPower2(Integral n)
{
    /*! 检查点数是否为 2 的幂形式 */
    bool isPower2 = false;
    if ( n != 0 && (n & (n-1)) == 0)
    {
        isPower2 = true;
    }
    return isPower2;
}
